package com.qiangesoft.ratelimit.core.semaphore;

import com.qiangesoft.ratelimit.core.LimitKeyBuilder;
import com.qiangesoft.ratelimit.core.LimitType;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;

/**
 * semaphore限流切面
 *
 * @author qiangesoft
 * @date 2024-03-21
 */
@Slf4j
@Aspect
@Component
public class SemaphoreLimitAspect {

    private static final Map<String, Semaphore> RATE_HOLDER = new ConcurrentHashMap<>();

    @Pointcut("@annotation(com.qiangesoft.ratelimit.core.semaphore.SemaphoreRateLimit)")
    public void pointcut() {
    }

    @Around(value = "pointcut()")
    public Object doAround(ProceedingJoinPoint joinPoint) throws Throwable {
        // 1.从方法获取注解
        SemaphoreRateLimit semaphoreRateLimit = this.getAnnotation(joinPoint);
        if (Objects.isNull(semaphoreRateLimit)) {
            return joinPoint.proceed();
        }

        // 2.限流逻辑
        LimitType limitType = semaphoreRateLimit.limitType();
        String key = semaphoreRateLimit.key();
        int count = semaphoreRateLimit.count();
        String limitKey = LimitKeyBuilder.build(limitType, key, joinPoint);
        Semaphore semaphore = getSemaphore(limitKey, count);
        try {
            if (!semaphore.tryAcquire()) {
                throw new RuntimeException("访问太频繁，请稍后再试！");
            }
            return joinPoint.proceed();
        } finally {
            semaphore.release();
        }
    }

    /**
     * 获取RateLimiter
     *
     * @param limitKey
     * @param permits
     * @return
     */
    public synchronized static Semaphore getSemaphore(String limitKey, int permits) {
        Semaphore semaphore = RATE_HOLDER.get(limitKey);
        if (semaphore != null) {
            return semaphore;
        }

        semaphore = new Semaphore(permits);
        RATE_HOLDER.put(limitKey, semaphore);
        return semaphore;
    }

    /**
     * 获取注解
     *
     * @param joinPoint
     * @return
     */
    private SemaphoreRateLimit getAnnotation(ProceedingJoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        return AnnotationUtils.findAnnotation(signature.getMethod(), SemaphoreRateLimit.class);
    }

}